{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 几个重要的点:\n",
    "    \n",
    "    Fixed Q-target(定时共享参数)\n",
    "    \n",
    "    Experience Repaly Memory(numpy)\n",
    "    \n",
    "    Gradient clipping(增加稳定性)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import gym\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from IPython.display import clear_output\n",
    "from torch.nn.utils import clip_grad_norm_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "    def __init__(self, column, max_size, batch_size):\n",
    "        self.current_state = np.zeros((max_size, column), dtype=np.float32)\n",
    "        self.next_state = np.zeros((max_size, column), dtype=np.float32)\n",
    "        self.action = np.zeros(max_size, dtype=np.float32)\n",
    "        self.reward = np.zeros(max_size, dtype=np.float32)\n",
    "        self.done = np.zeros(max_size,dtype=np.float32)\n",
    "        self.max_size, self.batch_size = max_size, batch_size\n",
    "        self.size, self.current_index = 0, 0\n",
    "    \n",
    "    def store(self, current_state, action, next_state, reward, done):\n",
    "        self.current_state[self.current_index] = current_state\n",
    "        self.action[self.current_index] = action\n",
    "        self.next_state[self.current_index] = next_state\n",
    "        self.reward[self.current_index] = reward\n",
    "        self.done[self.current_index] = done\n",
    "        self.current_index = (self.current_index + 1) % self.max_size\n",
    "        self.size = min(self.size + 1, self.max_size)\n",
    "    \n",
    "    def sample_batch(self):\n",
    "        ptr = np.random.choice(self.size, self.batch_size)\n",
    "        return dict(current_state=self.current_state[ptr],\n",
    "                    next_state=self.next_state[ptr],\n",
    "                    action=self.action[ptr],\n",
    "                    reward=self.reward[ptr],\n",
    "                    done=self.done[ptr]\n",
    "        )\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Network(nn.Module):\n",
    "    def __init__(self, in_dim, out_dim):\n",
    "        super(Network, self).__init__()\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            nn.Linear(in_dim, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, out_dim)\n",
    "        )\n",
    "    def forward(self,x):\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_epsilon = 0.05\n",
    "max_epsilon = 1\n",
    "epsilon_decay = 80\n",
    "epsilon_episode = lambda episode : min_epsilon + np.exp(-episode / epsilon_decay)*0.95\n",
    "plt.plot([epsilon_episode(x) for x in range(600)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v1')\n",
    "\n",
    "seed = 777\n",
    "\n",
    "def seed_torch(seed):\n",
    "    torch.manual_seed(seed)\n",
    "    if torch.backends.cudnn.enabled:\n",
    "        torch.backends.cudnn.benchmark = False\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "\n",
    "batch_size = 32\n",
    "max_size = 1000\n",
    "memory = ReplayBuffer(env.observation_space.shape[0], max_size, batch_size)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "network = Network(env.observation_space.shape[0],env.action_space.n).to(device)\n",
    "target_network = Network(env.observation_space.shape[0],env.action_space.n).to(device)\n",
    "target_network.load_state_dict(network.state_dict())\n",
    "target_network.eval()#把train设置为false，默认为true。可通过target_network.train()将train设置回true\n",
    "\n",
    "optimizer = optim.Adam(network.parameters(), lr=1e-2)\n",
    "\n",
    "gamma = 0.99\n",
    "target_update = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_action(episode, state):\n",
    "    if np.random.random_sample() > epsilon_episode(episode):\n",
    "        selected_action = network(torch.FloatTensor(state).to(device)).argmax().detach().cpu().numpy()\n",
    "    else:\n",
    "        selected_action = env.action_space.sample()\n",
    "    return selected_action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    samples = memory.sample_batch()\n",
    "    state = torch.FloatTensor(samples[\"current_state\"]).to(device).to(device)\n",
    "    next_state = torch.FloatTensor(samples[\"next_state\"]).to(device)\n",
    "    action = torch.LongTensor(samples[\"action\"].reshape(-1, 1)).to(device)\n",
    "    reward = torch.FloatTensor(samples[\"reward\"].reshape(-1, 1)).to(device)\n",
    "    done = torch.FloatTensor(samples[\"done\"].reshape(-1, 1)).to(device)\n",
    "    \n",
    "    current_Q_value = network(state).gather(1, action)\n",
    "#     gather : 对源多维数组的信息从新提取,第一个参数代表跨行(0)&跨列(1)\n",
    "#     tensor[[1,2],    index[[0,0],\n",
    "#            [3,4]]          [1,0]]\n",
    "#     dim = 0 return [[1,2],   dim = 1 return [[1,1],\n",
    "#                     [3,2]]                    [4,3]]\n",
    "    next_Q_value = target_network(next_state).max(dim=1,keepdim=True)[0].detach()\n",
    "#     .max()\n",
    "#         keepdim=True,保持原来的形状，如:(2,3) \\to (2,1)\n",
    "#     返回值有values 和 indices 所以要有[0]操作\n",
    "#     detach()\n",
    "#         共享tensor() not require grad\n",
    "    target = (reward + gamma*next_Q_value*(1 - done)).to(device)\n",
    "    loss = ((target - current_Q_value).pow(2)).mean()\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    clip_grad_norm_(network.parameters(),1.0,norm_type=1) # Gradient clipping(增加稳定性)\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_reward(frame_idx, rewards):\n",
    "    clear_output(True)\n",
    "    plt.figure(figsize=(20,5))\n",
    "    plt.title('frame %s. mean_reward: %s'%(frame_idx, np.mean(rewards[-10:])))\n",
    "    plt.plot(rewards)\n",
    "    plt.xlabel('epoch')\n",
    "    plt.ylabel('reward')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_rewards = []\n",
    "update_count = 0\n",
    "\n",
    "for episode in range(600):\n",
    "    state, rewards = env.reset(), 0\n",
    "    for i in range(10000):\n",
    "        action = select_action(episode, state)\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        memory.store(state, action, next_state, reward, done)\n",
    "        state = next_state\n",
    "        rewards += reward\n",
    "        if done:\n",
    "            break\n",
    "        if len(memory) > batch_size:\n",
    "            train()\n",
    "            update_count += 1\n",
    "            if update_count % target_update == 0:\n",
    "                target_network.load_state_dict(network.state_dict())     \n",
    "    all_rewards.append(rewards)\n",
    "    plot_reward(episode,all_rewards)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
